"""
This evaluation script follows 3DDFA and 3DDFA_V2
https://github.com/cleardusk/3DDFA
https://github.com/cleardusk/3DDFA_V2
"""

import os.path as osp
import numpy as np
from math import sqrt
from math import cos, atan2, asin
from mind3d.utils.synergynet_util import _load
from mind3d.utils.synergynet_util import ParamsPack

param_pack = ParamsPack()

d = '/data1/hujingsong/synergynet/aflw2000_data/eval'
yaws_list = _load(osp.join(d, 'AFLW2000-3D.pose.npy'))
# origin
pts68_all_ori = _load(osp.join(d, 'AFLW2000-3D.pts68.npy'))
# reannonated
pts68_all_re = _load(osp.join(d, 'AFLW2000-3D-Reannotated.pts68.npy'))
roi_boxs = _load(osp.join(d, 'AFLW2000-3D_crop.roi_box.npy'))


def ana(nme_list):
    yaw_list_abs = np.abs(yaws_list)
    ind_yaw_1 = yaw_list_abs <= 30
    ind_yaw_2 = np.bitwise_and(yaw_list_abs > 30, yaw_list_abs <= 60)
    ind_yaw_3 = yaw_list_abs > 60

    nme_1 = nme_list[ind_yaw_1]
    nme_2 = nme_list[ind_yaw_2]
    nme_3 = nme_list[ind_yaw_3]

    mean_nme_1 = np.mean(nme_1) * 100
    mean_nme_2 = np.mean(nme_2) * 100
    mean_nme_3 = np.mean(nme_3) * 100

    std_nme_1 = np.std(nme_1) * 100
    std_nme_2 = np.std(nme_2) * 100
    std_nme_3 = np.std(nme_3) * 100

    mean_all = [mean_nme_1, mean_nme_2, mean_nme_3]
    mean = np.mean(mean_all)
    std = np.std(mean_all)

    s1 = '[ 0, 30]\tMean: \x1b[32m{:.3f}\x1b[0m, Std: {:.3f}'.format(mean_nme_1, std_nme_1)
    s2 = '[30, 60]\tMean: \x1b[32m{:.3f}\x1b[0m, Std: {:.3f}'.format(mean_nme_2, std_nme_2)
    s3 = '[60, 90]\tMean: \x1b[32m{:.3f}\x1b[0m, Std: {:.3f}'.format(mean_nme_3, std_nme_3)
    s5 = '[ 0, 90]\tMean: \x1b[31m{:.3f}\x1b[0m, Std: \x1b[31m{:.3f}\x1b[0m'.format(mean, std)

    s = '\n'.join([s1, s2, s3, s5])

    print(s)

    return mean_nme_1, mean_nme_2, mean_nme_3, mean, std


def ana_msg(nme_list):
    leng = nme_list.shape[0]
    yaw_list_abs = np.abs(yaws_list)[:leng]
    ind_yaw_1 = yaw_list_abs <= 30
    ind_yaw_2 = np.bitwise_and(yaw_list_abs > 30, yaw_list_abs <= 60)
    ind_yaw_3 = yaw_list_abs > 60

    nme_1 = nme_list[ind_yaw_1]
    nme_2 = nme_list[ind_yaw_2]
    nme_3 = nme_list[ind_yaw_3]

    mean_nme_1 = np.mean(nme_1) * 100
    mean_nme_2 = np.mean(nme_2) * 100
    mean_nme_3 = np.mean(nme_3) * 100

    std_nme_1 = np.std(nme_1) * 100
    std_nme_2 = np.std(nme_2) * 100
    std_nme_3 = np.std(nme_3) * 100

    mean_all = [mean_nme_1, mean_nme_2, mean_nme_3]
    mean = np.mean(mean_all)
    std = np.std(mean_all)

    s0 = '\nFacial Alignment on AFLW2000-3D (NME):'
    s1 = '[ 0, 30]\tMean: {:.3f}, Std: {:.3f}'.format(mean_nme_1, std_nme_1)
    s2 = '[30, 60]\tMean: {:.3f}, Std: {:.3f}'.format(mean_nme_2, std_nme_2)
    s3 = '[60, 90]\tMean: {:.3f}, Std: {:.3f}'.format(mean_nme_3, std_nme_3)
    s4 = '[ 0, 90]\tMean: {:.3f}, Std: {:.3f}'.format(mean, std)

    s = '\n'.join([s0, s1, s2, s3, s4])

    return s


def convert_to_ori(lms, i):
    std_size = 120
    sx, sy, ex, ey = roi_boxs[i]
    scale_x = (ex - sx) / std_size
    scale_y = (ey - sy) / std_size
    lms[0, :] = lms[0, :] * scale_x + sx
    lms[1, :] = lms[1, :] * scale_y + sy
    return lms


def convert_to_crop(lms, i):
    std_size = 120
    sx, sy, ex, ey = roi_boxs[i]
    scale_x = (ex - sx) / std_size
    scale_y = (ey - sy) / std_size
    lms[0, :] = (lms[0, :] - sx) / scale_x
    lms[1, :] = (lms[1, :] - sy) / scale_y
    return lms


def calc_nme(pts68_fit_all, option='ori'):
    if option == 'ori':
        pts68_all = pts68_all_ori
    elif option == 're':
        pts68_all = pts68_all_re
    std_size = 120

    nme_list = []
    length_list = []
    for i in range(len(roi_boxs)):
        pts68_fit = pts68_fit_all[i]
        pts68_gt = pts68_all[i]

        sx, sy, ex, ey = roi_boxs[i]
        scale_x = (ex - sx) / std_size
        scale_y = (ey - sy) / std_size
        pts68_fit[0, :] = pts68_fit[0, :] * scale_x + sx
        pts68_fit[1, :] = pts68_fit[1, :] * scale_y + sy

        # build bbox
        minx, maxx = np.min(pts68_gt[0, :]), np.max(pts68_gt[0, :])
        miny, maxy = np.min(pts68_gt[1, :]), np.max(pts68_gt[1, :])
        llength = sqrt((maxx - minx) * (maxy - miny))
        length_list.append(llength)

        dis = pts68_fit - pts68_gt[:2, :]
        dis = np.sqrt(np.sum(np.power(dis, 2), 0))
        dis = np.mean(dis)
        nme = dis / llength
        nme_list.append(nme)

    nme_list = np.array(nme_list, dtype=np.float32)
    return nme_list


def p2srt(p):
    '''decomposing camera matrix P'''
    t3d = p[:, 3]
    r1 = p[0:1, :3]
    r2 = p[1:2, :3]
    s = (np.linalg.norm(r1) + np.linalg.norm(r2)) / 2.0
    r1 = r1 / np.linalg.norm(r1)
    r2 = r2 / np.linalg.norm(r2)
    r3 = np.cross(r1, r2)
    r = np.concatenate((r1, r2, r3), 0)
    return s, r, t3d


def matrix2angle(r):
    '''convert matrix to angle'''
    if r[2, 0] != 1 and r[2, 0] != -1:
        x = asin(r[2, 0])
        y = atan2(r[1, 2] / cos(x), r[2, 2] / cos(x))
        z = atan2(r[0, 1] / cos(x), r[0, 0] / cos(x))

    else:  # Gimbal lock
        z = 0  # can be anything
        if r[2, 0] == -1:
            x = np.pi / 2
            y = z + atan2(r[0, 1], r[0, 2])
        else:
            x = -np.pi / 2
            y = -z + atan2(-r[0, 1], -r[0, 2])

    rx, ry, rz = x * 180 / np.pi, y * 180 / np.pi, z * 180 / np.pi

    return [rx, ry, rz]


def parse_pose(param):
    '''parse parameters into pose'''
    if len(param) == 62:
        param = param * param_pack.param_std[:62] + param_pack.param_mean[:62]
    else:
        param = param * param_pack.param_std + param_pack.param_mean
    ps = param[:12].reshape(3, -1)  # camera matrix
    _, r, t3d = p2srt(ps)
    p = np.concatenate((r, t3d.reshape(3, -1)), axis=1)  # without scale
    pose = matrix2angle(r)  # yaw, pitch, roll
    return p, pose


def parsing(param):
    """
    parsing function
    """
    p_ = param[:, :12].reshape(-1, 3, 4)
    p = p_[:, :, :3]
    offset = p_[:, :, -1].reshape(-1, 3, 1)
    alpha_shp = param[:, 12:52].reshape(-1, 40, 1)
    alpha_exp = param[:, 52:62].reshape(-1, 10, 1)
    return p, offset, alpha_shp, alpha_exp


if __name__ == '__main__':
    pass
